import whisper
import opencc
import numpy as np
from flask import Flask, request, jsonify
import os
from werkzeug.utils import secure_filename
from concurrent.futures import ThreadPoolExecutor
import queue
import threading
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = './uploads'
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# 初始化模型和转换器
model = whisper.load_model("large-v3")
cc = opencc.OpenCC('t2s.json')

# 并发配置
MAX_CONCURRENT_TASKS = 1  # 保持为1最稳定
TASK_QUEUE_MAX_SIZE = 10

# 任务队列和线程池
task_queue = queue.Queue(maxsize=TASK_QUEUE_MAX_SIZE)
executor = ThreadPoolExecutor(max_workers=MAX_CONCURRENT_TASKS)
model_lock = threading.Lock()

def validate_whisper_result(result):
    """验证Whisper返回结果的格式"""
    if not isinstance(result, dict):
        raise ValueError("Result is not a dictionary")
    if 'text' not in result:
        raise ValueError("Missing 'text' field in result")
    if 'segments' not in result:
        raise ValueError("Missing 'segments' field in result")
    return True

def process_audio(audio_path):
    """处理音频的核心函数"""
    with model_lock:
        try:
            logger.info(f"Processing audio: {audio_path}")
            
            # 1. 执行转写
            result = model.transcribe(audio_path, word_timestamps=True)
            if result is None:
                raise ValueError("Whisper returned None")
                
            # 2. 验证结果格式
            validate_whisper_result(result)
            
            # 3. 处理文本
            result['text'] = cc.convert(result.get('text', ''))
            
            # 4. 处理分段
            segments = []
            for segment in result.get('segments', []):
                if not isinstance(segment, dict):
                    continue
                    
                processed_segment = {
                    'start': f"{np.round(segment.get('start', 0), 2)}",
                    'end': f"{np.round(segment.get('end', 0), 2)}",
                    'text': cc.convert(segment.get('text', ''))
                }
                segments.append(processed_segment)
            
            # 5. 重建结果字典
            final_result = {
                'text': result['text'],
                'segments': segments
            }
            
            return final_result
            
        except Exception as e:
            logger.error(f"Error in process_audio: {str(e)}")
            raise

def worker():
    """工作线程"""
    while True:
        try:
            task = task_queue.get()
            if task is None:
                break
                
            future, save_path, callback = task
            try:
                result = process_audio(save_path)
                future.set_result(result)
            except Exception as e:
                future.set_exception(e)
            finally:
                if os.path.exists(save_path):
                    try:
                        os.remove(save_path)
                    except Exception as e:
                        logger.warning(f"Failed to remove {save_path}: {str(e)}")
                task_queue.task_done()
                if callback:
                    callback()
                    
        except Exception as e:
            logger.error(f"Worker error: {str(e)}")

worker_thread = threading.Thread(target=worker)
worker_thread.daemon = True
worker_thread.start()

@app.route('/transcribe', methods=['POST'])
def transcribe_api():
    if 'file' not in request.files:
        return jsonify({'code': 400, 'error': 'No file uploaded'})
    
    file = request.files['file']
    if not file.filename:
        return jsonify({'code': 400, 'error': 'Empty filename'})
    
    save_path = None
    try:
        if task_queue.qsize() >= TASK_QUEUE_MAX_SIZE:
            return jsonify({'code': 503, 'error': 'Server busy'})
        
        filename = secure_filename(file.filename)
        save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(save_path)
        
        future = executor.submit(lambda: None)
        future._state = 'PENDING'
        
        task_queue.put((future, save_path, None))
        
        result = future.result()
        if not isinstance(result, dict):
            raise ValueError("Invalid result format")
            
        return jsonify({
            'code': 0,
            'text': result.get('text', ''),
            'segments': result.get('segments', [])
        })
        
    except Exception as e:
        logger.error(f"API error: {str(e)}")
        if save_path and os.path.exists(save_path):
            try:
                os.remove(save_path)
            except:
                pass
        return jsonify({
            'code': 500,
            'error': str(e),
            'message': 'Internal server error'
        })

@app.route('/queue-status', methods=['GET'])
def queue_status():
    return jsonify({
        'code': 0,
        'queue_size': task_queue.qsize(),
        'max_concurrent': MAX_CONCURRENT_TASKS
    })

if __name__ == '__main__':
    try:
        app.run(host='0.0.0.0', port=5000, threaded=True)
    finally:
        task_queue.put(None)
        worker_thread.join()